import numpy
import operator
import gdal

class Raster:
    def __init__(self, data, transform, nodata=None):
        self.data = data
        self.transform = transform
        self.nodata = nodata
        if len(self.data.shape)==2:
            self.bands = 1
            self.rows = self.data.shape[0]
            self.cols = self.data.shape[1]
        else:
            self.bands = self.data.shape[0]
            self.rows = self.data.shape[1]
            self.cols = self.data.shape[2]

    def saveFile(self, filePath, format="GTiff"):
        driver = gdal.GetDriverByName(format)
        dst_ds = driver.Create(filePath,self.cols,self.rows,self.bands,gdal.GDT_Float32)
        dst_ds.SetGeoTransform(self.transform)
        try:
            # this seems a fair test for if we have to convert to Numeric arrays:
            import gdalnumeric
            import Numeric
            data = Numeric.array(self.data)
        except ImportError:
            data = self.data
        if self.bands == 1:
            dst_ds.GetRasterBand(1).WriteArray(data)
        else:
            for band in range(self.bands):
                dst_ds.GetRasterBand(band+1).WriteArray(data[band])
        dst_ds = None

    def info(self):
        print "Shape: ",self.data.shape
        print "Transform: ",self.transform
        print "Nodata: ",self.nodata
        print "Bands: ",self.bands
        print "Rows: %s Cols: %s\n" % (self.rows,self.cols)

    def __repr__(self):
        return "Raster: "+str(self.data.shape)
###
### maths methods
###
    def __math__(self,other,op):
        if isinstance(other,Raster):
            print "raster+raster"
            print other,self
            try:
                r = Raster(op(self.data,other.data),self.transform,self.nodata)
            except:
                print "raster math error"
        else:
            r = Raster(op(self.data,other),self.transform,self.nodata)
        return r

    def __add__(self,other):
        return self.__math__(other,operator.add)
    def __radd__(self,other):
        return self + other

    def __sub__(self,other):
        return self.__math__(other,operator.sub)
    def __rsub__(self,other):
        return self + other

    def __mul__(self,other):
        return self.__math__(other,operator.mul)
    def __rmul__(self,other):
        return self * other

    def __div__(self,other):
        return self.__math__(other,operator.div)
    def __rdiv__(self,other):
        r = Raster(other/self.data, self.transform, self.nodata)
        return r

    def __pow__(self,other):
        return self.__math__(other,operator.pow)
    def __rpow__(self,other):
        r = Raster( other ** self.data, self.transform, self.nodata)
        return r

    def __abs__(self):
        r = Raster(abs(self.data), self.transform, self.nodata)
        return r
    
    def __neg__(self):
        r = Raster(-self.data,self.transform,self.nodata)
        return r
        
def fromFile(filePath):
    gd = gdal.Open(str(filePath))
    array = gd.ReadAsArray()
    transform = gd.GetGeoTransform()
    nodata=[]
    for ib in range(1,gd.RasterCount):
        nodata.append(gd.GetRasterBand(ib).GetNoDataValue())
    return Raster(array, transform, nodata)

def fromSimpleRaster(sr):
    """ simpleRaster is a dict with:
    columns, rows, 
    xleft ytop
    cellHeight cellWidth
    startValue incrementValue
    """
    ncells = sr['columns'] * sr['rows']
    lastValue = sr['startValue'] + sr['incrementValue'] * (ncells-1)
    data = numpy.linspace(sr['startValue'],lastValue,ncells).reshape(sr['columns'],sr['rows'])
    transform = [sr['xleft'],sr['cellWidth'],0,sr['ytop'],0,sr['cellHeight']]
    nodata = None
    return Raster(data,transform,nodata)
    

import sys
if __name__ == "__main__":
    if len(sys.argv)>1:
        fn = sys.argv[1]
        r = fromFile(fn)
        r.info()
        r.saveFile("output.tif")

    t = dict(columns=5, rows=3, startValue=23, incrementValue=4.1, xleft=-3.1, ytop=54, cellWidth=0.1, cellHeight=0.12)
    tr = fromSimpleRaster(t)
    tr.info()
    tr.saveFile("output2.tif")

    t1 = fromSimpleRaster(dict(columns=2, rows=2, startValue=1, incrementValue=1, xleft=0, ytop=0, cellWidth=1, cellHeight=1))
    t2 = fromSimpleRaster(dict(columns=2, rows=2, startValue=2, incrementValue=1.5, xleft=0, ytop=0, cellWidth=1, cellHeight=1))
    
    t3 = t1 + t2
    t4 = t1 - t2

    print t3.data
    print t4.data

    t4.saveFile("t4.tif")
